package com.ktjiaoyu.travelaiagent.rag;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 旅游文档读取器
 */
@Component
@Slf4j
public class TravelMarkdownDocumentReader implements DocumentReader {

    @Autowired
    private ResourcePatternResolver resourcePatternResolver;
    // 指定Markdown文档的路径模式，默认读取classpath*/doc/*.md
    @Value("${travel.document.markdown.path.pattern:classpath*:/doc/*.md}")
    private String pathPattern;

    @Override
    public List<Document> get() {
        List<Document> documents = new ArrayList<>();
        try {
            // 使用ResourcePatternResolver获取所有匹配的Markdown文档
            Resource[] markdownResources = resourcePatternResolver.getResources(pathPattern);
            for (Resource resource : markdownResources) {
                // 构建读取器配置对象，可以根据需要配置不同的参数
                MarkdownDocumentReaderConfig readerConfig = MarkdownDocumentReaderConfig.builder()
                        .withHorizontalRuleCreateDocument(true)
                        .withIncludeCodeBlock(false)
                        .withIncludeBlockquote(false)
                        .withAdditionalMetadata("filename", resource.getFilename())
                        .build();
                // 创建读取器
                MarkdownDocumentReader reader = new MarkdownDocumentReader(resource, readerConfig);
                // 读取文档并添加到列表中
                documents.addAll(reader.get());
            }
        } catch (IOException e) {
            // 这里可以根据实际需要进行日志记录或异常处理
            log.error("旅游信息Markdown文档读取失败。", e);
        }
        return documents;
    }
}
